import torch
import torch.nn as nn


class PromptTuning(nn.Module):
    def __init__(self, model, prompt_length=10):
        super().__init__()
        self.model = model
        self.prompt = nn.Parameter(torch.randn(prompt_length, model.config.hidden_size))

    def forward(self, input_ids):
        # 获取输入嵌入
        inputs_embeds = self.model.get_input_embeddings()(input_ids)

        # 添加提示向量
        prompt_embeds = self.prompt.repeat(inputs_embeds.size(0), 1, 1)
        inputs_embeds = torch.cat([prompt_embeds, inputs_embeds], dim=1)

        # 前向传播（冻结原始模型）
        with torch.no_grad():
            outputs = self.model(inputs_embeds=inputs_embeds)
        return outputs